
from scripts.model import MySegNet
import pytorch_lightning

def train(lightninig_param, train_dl, val_dl):
    net = MySegNet()
    trainer = pytorch_lightning.Trainer(max_epochs=lightninig_param['max_epochs'], 
                                        default_root_dir=lightninig_param['default_root_dir'],
                                        check_val_every_n_epoch=lightninig_param['check_val_every_n_epoch'],
                                        devices=1, num_nodes=1, log_every_n_steps=3, num_sanity_val_steps=0)
    trainer.fit(model=net, train_dataloaders=train_dl, val_dataloaders=val_dl)


def evaluate(lightninig_param, ckpt_file, val_dl):
    net = MySegNet()
    trainer = pytorch_lightning.Trainer(default_root_dir=lightninig_param['default_root_dir'],
                                        devices=1, num_nodes=1)
    trainer.validate(model=net, dataloaders=val_dl, ckpt_path=ckpt_file)
